/*
 * Copyright 2007 Brian Watkins
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package bw.net;

import com.jcraft.jsch.*;
import java.io.*;
import java.util.*;
import java.util.zip.*;
import bw.util.*;

public class SCPClient
{
    private String _host = null;
    private String _user = null;
    private String _password = null;
    private JSch _jsch = null;

    private Session _session = null;
    private Channel _channel = null;

    public SCPClient ( String host, String user, String pass )
    {
        _host = host;
        _user = user;
        _password = pass;

        _jsch=new JSch();

        //turn off strict host key checking. probably should get it to look
        //at a known hosts file instead
        _jsch.setConfig( "StrictHostKeyChecking", "no" );

    }

    public void uploadFile ( File f, String remoteFile )
        throws Exception
    {
        uploadBytesTo( new FileInputStream( f ), f.length(), remoteFile );
    }

    public void uploadObjectTo ( Object obj, String remoteFile )
        throws Exception
    {
        //get the bytes. gzip it to save spacee
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        GZIPOutputStream gzout = new GZIPOutputStream( bos );
        ObjectOutputStream oos = new ObjectOutputStream( gzout );
        oos.writeObject( obj );
        oos.close();

        byte [] bytes = bos.toByteArray();

        uploadBytesTo( new ByteArrayInputStream( bytes ), bytes.length, remoteFile );
    }

    public void uploadBytesTo ( InputStream bis, long length, String remoteFile )
        throws Exception
    {
        Channel channel = openChannel( "scp -p -t " + remoteFile );

        // get I/O streams for remote scp
        OutputStream out = channel.getOutputStream();
        InputStream in = channel.getInputStream();

        channel.connect();

        if ( checkAck( in ) !=0 )
        {
            Log.getInstance().write( "Bad problem with SCP" );
        }
        
        // send "C0644 filesize filename", where filename should not include '/'
        String command = "C0644 "+ length + " wd-osFile.obj\n";

        out.write( command.getBytes() ); 
        out.flush();
        if ( checkAck( in ) !=0 )
        {
            Log.getInstance().write( "Problem sending scp command!" );
        }
            
        // send bytes
        byte[] buf=new byte[1024];
        while ( true )
        {
            int len = bis.read( buf, 0, buf.length );
            if ( len <= 0 ) break;
            out.write( buf, 0, len ); //out.flush();
        }
        bis.close();
        bis = null;
        
        // send '\0'
        sendReturn( buf, out );

        if ( checkAck( in ) !=0 ) 
        {
            System.out.println( "Problem sending bytes!" );
        }
        out.close();

        closeChannel( channel );
    }

    public Object downloadObjectFrom ( String remoteFile )
        throws Exception
    {
        byte [] bytes = downloadBytesFrom( remoteFile );
        
        //objects are gzipped
        GZIPInputStream gzin = new GZIPInputStream( new ByteArrayInputStream( bytes ) );

        ObjectInputStream ois = new ObjectInputStream( gzin );
        Object obj = ois.readObject();

        ois.close();

        return obj;
    }

    public byte [] downloadBytesFrom ( String remoteFile )
        throws Exception
    {
        Channel channel = openChannel( "scp -f " + remoteFile );
        
        // get I/O streams for remote scp
        OutputStream out = channel.getOutputStream();
        InputStream in = channel.getInputStream();

        channel.connect();

        byte[] buf = new byte[1024];

        // send '\0'
        sendReturn( buf, out );

        ByteArrayOutputStream bos = new ByteArrayOutputStream();

        //REVISIT: this will NOT work for copying multiple files
        while ( true )
        {
            int c = checkAck( in );
            if( c != 'C' )
            {
                break;
            }

            // read '0644 '
            in.read( buf, 0, 5 );
            
            long filesize = 0L;
            while( true )
            {
                if ( in.read( buf, 0, 1 ) < 0 )
                {
                    // error
                    break; 
                }
                if ( buf[0] == ' ' ) break;
                filesize = filesize * 10L + (long) ( buf[0] - '0' );
            }

            String file = null;
            for ( int i=0;; i++) 
            {
                in.read ( buf, i, 1 );
                if ( buf[i] == (byte) 0x0a )
                {
                    file = new String( buf, 0, i );
                    break;
                }
            }

            //System.out.println("filesize="+filesize+", file="+file);

            // send '\0'
            sendReturn( buf, out );
            
            // read a content of lfile
            int foo;
            while ( true )
            {
                if ( buf.length < filesize )
                {
                    foo = buf.length;
                }
                else
                {
                    foo = (int) filesize;
                }
                foo = in.read( buf, 0, foo );
                if ( foo < 0 )
                {
                    // error 
                    break;
                }
                bos.write( buf, 0, foo );
                filesize -= foo;
                if ( filesize == 0L ) break;
            }
            bos.close();
            
            if ( checkAck( in ) !=0 )
            {
                Log.getInstance().write( "Problem reading file!" );
            }

            // send '\0'
            sendReturn( buf, out );
        }
        
        closeChannel( channel );

        return bos.toByteArray();
    }

    private Channel openChannel ( String command )
        throws Exception
    {
        _session = _jsch.getSession( _user, _host, 22 );
        _session.setPassword( _password );
        _session.connect();

        Channel channel = _session.openChannel( "exec" );
        ( (ChannelExec) channel ).setCommand( command );

        return channel;
    }

    private void closeChannel ( Channel channel )
    {
        channel.disconnect();
        _session.disconnect();
        _session = null;
    }

    private void sendReturn ( byte [] buf, OutputStream out )
        throws Exception
    {
        buf[0] = 0;
        out.write( buf, 0, 1 ); 
        out.flush();
    }

    private int checkAck ( InputStream in ) 
        throws IOException
    {
        int b = in.read();
        // b may be 0 for success,
        //          1 for error,
        //          2 for fatal error,
        //          -1

        if( b == 0 ) return b;
        if ( b == -1 ) return b;
        
        if ( b == 1 || b == 2 )
        {
            StringBuffer sb = new StringBuffer();
            int c;
            do 
            {
                c = in.read();
                sb.append( (char) c );
            }
            while ( c != '\n' );
            if ( b == 1 ) 
            { // error
                Log.getInstance().write( sb.toString() );
            }
            if ( b == 2 )
            { // fatal error
                Log.getInstance().write( sb.toString() );
            }
        }
        
        return b;
    }

    public static void main ( String [] args )
    {
        try
        {
            if ( args.length != 4 )
            {
                System.out.println( "SCPClient [host] [user] [password] [rootDir]" );
                System.exit( 0 );
            }

            String host = args[0];
            String user = args[1];
            String password = args[2];
            String root = args[3];

            SCPClient client = new SCPClient( host, user, password );

            Vector v = new Vector();
            v.addElement( "one" );
            v.addElement( "two" );
            v.addElement( "three" );

            String remoteFile = root + "/wd-1.obj";

            System.out.print( "Uploading object ... " );
            client.uploadObjectTo( v, remoteFile );
            System.out.println( "Done!" );

            System.out.print( "Downloading object ... " );
            Vector g = (Vector) client.downloadObjectFrom( remoteFile );
            System.out.println( "Done!" );

            for ( Enumeration e = g.elements(); e.hasMoreElements(); )
            {
                System.out.println( (String) e.nextElement() );
            }
        }
        catch ( Exception ex )
        {
            ex.printStackTrace();
        }
    }
}